import sys

from PyQt6.QtCore import Qt
from PyQt6.QtWidgets import QApplication, QMessageBox
from ui import MedicalDetectionUI
from worker import DetectionWorker
from utils import get_model_path
from PyQt6.QtGui import QImage, QPixmap

def main():
    app = QApplication(sys.argv)
    window = MedicalDetectionUI()
    window.show()

    # 将检测按钮绑定线程
    def run_detection():
        if not window.current_image_path:
            QMessageBox.warning(window, "警告", "请先上传医学影像!")
            return

        window.upload_btn.setEnabled(False)
        window.detect_btn.setEnabled(False)
        window.detect_btn.setText("检测中...")
        window.progress_bar.setVisible(True)
        window.progress_bar.setValue(0)

        model_path = get_model_path()
        print(f"模型路径: {model_path}")  # 添加调试信息
        
        # 检查模型文件是否存在
        import os
        if not os.path.exists(model_path):
            QMessageBox.critical(window, "错误", f"模型文件不存在: {model_path}")
            window.reset_detection_ui()
            return

        worker = DetectionWorker(model_path, window.current_image_path)
        # 保存worker引用，防止被垃圾回收
        window.worker = worker

        def on_progress(value):
            if window.progress_bar:
                window.progress_bar.setValue(value)

        def on_finished(result):
            if window.progress_bar:
                window.progress_bar.setVisible(False)
            if window.upload_btn:
                window.upload_btn.setEnabled(True)
            if window.detect_btn:
                window.detect_btn.setEnabled(True)
                window.detect_btn.setText("病灶检测")

            if "error" in result:
                QMessageBox.critical(window, "错误", f"检测过程中发生错误:\n{result['error']}")
                return

            # 通过UI类方法显示结果
            window.show_detection_results(result)

            # 显示完成消息
            num_detections = result['num_detections']
            if num_detections > 0:
                QMessageBox.information(window, "检测完成", f"检测完成! 发现 {num_detections} 个病灶.")
            else:
                QMessageBox.information(window, "检测完成", "检测完成! 未发现明显病灶.")

        worker.progress.connect(on_progress)
        worker.finished.connect(on_finished)
        worker.start()

    window.detect_btn.clicked.connect(run_detection)
    
    # 保证程序正常运行
    exit_code = app.exec()
    sys.exit(exit_code)

if __name__ == "__main__":
    main()